/*
 * Project: MoveAce
 * File: moveace.c
 *
 * Copyright (C) 2009 Daniel Meekins
 * Contact: dmeekins - gmail
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

/*
 * This tool moves an Access Control Entry (ACE) within an Acces Control
 * List (ACL) for a file or directory on NTFS. It can also reorder the entries
 * in the list at once.
 *
 * Useful links for the win32 api used here:
 * http://www.devx.com/cplus/Article/16711/1954
 * http://msdn.microsoft.com/en-us/library/aa446595%28VS.85%29.aspx
 * http://msdn.microsoft.com/en-us/library/aa379608%28VS.85%29.aspx
 */

// NOTE: if this doesn't compile for some reason, try uncommenting the
// the following #define (0x0502 is for WinXP).
// (see http://msdn.microsoft.com/en-us/library/aa383745%28VS.85%29.aspx)
//#define _WIN32_WINNT 0x0502

#include <stdio.h>
#include <string.h>
#include <windows.h>
#include <winbase.h>
#include <sddl.h>

#define MODE_LIST    1
#define MODE_MOVE    2
#define MODE_REORDER 3

// Combined general permissions for simplicity
#define READ_PERMISSIONS (FILE_READ_DATA | FILE_READ_ATTRIBUTES)
#define WRITE_PERMISSIONS (FILE_WRITE_DATA | FILE_APPEND_DATA | FILE_WRITE_ATTRIBUTES | FILE_WRITE_EA)
#define EXECUTE_PERMISSIONS (FILE_READ_DATA | FILE_EXECUTE)

void print_ace_arr(ACCESS_ALLOWED_ACE **ace_arr, DWORD len);
BOOL move_ace(ACCESS_ALLOWED_ACE **ace_arr, int len, int src, int dst);
BOOL reorder_acl(ACCESS_ALLOWED_ACE **ace_arr, int len, int *idx_order, int idx_len);
BOOL isnum(char *str);
void print_error(LPWSTR func_name, DWORD error_code);
int usage(char *progname);

/*
 *
 */
int main(int argc, char *argv[])
{
    int i;
    DWORD dw;
    int mode = 0;
    LPWSTR filename = NULL;
    int *new_idx_list = NULL;
    int idx_list_len;

    SECURITY_DESCRIPTOR *sec_desc = NULL;
    SECURITY_INFORMATION sec_info = DACL_SECURITY_INFORMATION;
    DWORD size_needed;
    
    PACL dacl = NULL;
    BOOL dacl_present;
    BOOL dacl_defaulted;
    ACCESS_ALLOWED_ACE *ace = NULL;
    ACL_SIZE_INFORMATION acl_size;
    ACCESS_ALLOWED_ACE **ace_arr = NULL;

    if(argc < 2 || strcmp(argv[1], "/?") == 0 || strcmp(argv[1], "-?") == 0)
        return usage(argv[0]);

    // argv[1] should be file name
    // get the length needed
    i = MultiByteToWideChar(CP_ACP, 0, argv[1], -1, NULL, 0); 
    filename = (LPWSTR)malloc(i * sizeof(WCHAR));
    
    // copy arg to wide string
    MultiByteToWideChar(CP_ACP, 0, argv[1], -1, filename, i);
    mode = MODE_LIST;

    // parse the rest of the args
    if(argc > 2)
    {
        // allocate new index array based on number of indicies provided
        idx_list_len = argc - 3;
        new_idx_list = (int *)calloc(idx_list_len, sizeof(int));

        if(_stricmp(argv[2], "/m") == 0 || _stricmp(argv[2], "-m") == 0)
        {
            if(argc != 5 || !isnum(argv[3]) || !isnum(argv[4]))
                goto parse_error;

            new_idx_list[0] = atoi(argv[3]);
            new_idx_list[1] = atoi(argv[4]);

            mode = MODE_MOVE;
        }
        else if(_stricmp(argv[2], "/r") == 0 || _stricmp(argv[2], "-r") == 0)
        {
            if(argc < 4)
                goto parse_error;

            for(i = 3; i < argc; i++)
            {
                if(!isnum(argv[i]))
                    goto parse_error;
                new_idx_list[i - 3] = atoi(argv[i]);
            }

            mode = MODE_REORDER;
        }
        else
        {
            goto parse_error;
        }
    }

    wprintf(L"Processing file %s\n", filename);

    // first need to get the security descriptor from the file
    // call once to get the amount of memory needed
    GetFileSecurity(filename, sec_info, NULL, 0, &size_needed);
    if(GetLastError() != ERROR_INSUFFICIENT_BUFFER)
    {
        print_error(L"GetFileSecurity", GetLastError());
        goto done;
    }

    // allocate needed memory for security descriptor
    if((sec_desc = (SECURITY_DESCRIPTOR *)calloc(1, size_needed)) == NULL)
    {
        printf("ERROR: calloc() failed for sec_desc\n");
        goto done;
    }

    // get the security descriptor from the file
    if(!GetFileSecurity(filename, sec_info, sec_desc, size_needed, &size_needed))
    {
        print_error(L"GetFileSecurity", GetLastError());
        goto done;
    }

    // get the security descriptor DACL
    if(!GetSecurityDescriptorDacl(sec_desc, &dacl_present, &dacl, &dacl_defaulted))
    {
        print_error(L"GetSecurityDescriptorDacl", GetLastError());
        goto done;
    }
    else if(!dacl_present || dacl == NULL)
    {
        printf("No DACL found.\n");
        goto done;
    }

    ZeroMemory(&acl_size, sizeof(ACL_SIZE_INFORMATION));
    acl_size.AclBytesInUse = sizeof(ACL);

    // get the size information of the ACL for the list length
    if(!GetAclInformation(dacl, (LPVOID)&acl_size, sizeof(ACL_SIZE_INFORMATION), AclSizeInformation))
    {
        print_error(L"GetAclInformation", GetLastError());
        goto done;
    }

    // create array to store all ACEs in
    if((ace_arr = (ACCESS_ALLOWED_ACE **)calloc(acl_size.AceCount, sizeof(ACCESS_ALLOWED_ACE *))) == NULL)
    {
        printf("ERROR: calloc() failed\n");
        goto done;
    }

    // put copies of all the ACEs into the array
    for(dw = 0; dw < acl_size.AceCount; dw++)
    {
        if(!GetAce(dacl, dw, (LPVOID)&ace))
        {
            print_error(L"GetAce", GetLastError());
            goto done;
        }

        ace_arr[dw] = (ACCESS_ALLOWED_ACE *)calloc(1, ace->Header.AceSize);
        memcpy(ace_arr[dw], ace, ace->Header.AceSize);
    }

    // finally perform the operation based on the command's arguments
    switch(mode)
    {
        case MODE_LIST:
            print_ace_arr(ace_arr, acl_size.AceCount);
            break;

        case MODE_MOVE:
            if(!move_ace(ace_arr, (int)acl_size.AceCount, new_idx_list[0], new_idx_list[1]))
                goto done;
            //print_ace_arr(ace_arr, acl_size.AceCount); // debug
            break;

        case MODE_REORDER:
            if(!reorder_acl(ace_arr, (int)acl_size.AceCount, new_idx_list, idx_list_len))
                goto done;
            //print_ace_arr(ace_arr, acl_size.AceCount); // debug
            break;
    }

    if(mode == MODE_MOVE || mode == MODE_REORDER)
    {
        // remove all the entries
        for(dw = 0; dw < acl_size.AceCount; dw++)
        {
            if(!DeleteAce(dacl, 0))
            {
                print_error(L"DeleteAce", GetLastError());
                goto done;
            }
        }

        // add them back again in the new order
        for(dw = 0; dw < acl_size.AceCount; dw++)
        {
            if(!AddAce(dacl, ACL_REVISION, MAXDWORD, (LPVOID)ace_arr[dw], ace_arr[dw]->Header.AceSize))
            {
                print_error(L"AddAce", GetLastError());
                goto done;
            }
        }

        // write back security descriptor (includes new dacl)
        if(!SetFileSecurity(filename, sec_info, sec_desc))
        {
            print_error(L"SetFileSecurity", GetLastError());
            goto done;
        }
    }

    printf("Done.\n"); //yay

  done:
    if(filename)
        free(filename);
    if(new_idx_list)
        free(new_idx_list);
    if(ace_arr)
    {
        for(dw = 0; dw < acl_size.AceCount; dw++)
        {
            if(ace_arr[dw])
                free(ace_arr[dw]);
        }
        free(ace_arr);
    }
    if(sec_desc)
        free(sec_desc);
    return 0;

  parse_error:
    if(filename)
        free(filename);
    if(new_idx_list)
        free(new_idx_list);
    return usage(argv[0]);
}

/*
 * Prints the list of ACEs
 */
void print_ace_arr(ACCESS_ALLOWED_ACE **ace_arr, DWORD len)
{
    DWORD i;
    LPTSTR sid_str = NULL;
    
    for(i = 0; i < len; i++)
    {
        // print only the access_allowed or access_denied types
        if(ace_arr[i]->Header.AceType == ACCESS_ALLOWED_ACE_TYPE)
            printf("  [%d]: ALLOW, ", i);
        else if(ace_arr[i]->Header.AceType == ACCESS_DENIED_ACE_TYPE)
            printf("  [%d]: DENY,  ", i);
        else
            continue;
        
        // print the permissions
        printf("%08x [%c%c%c], ", ace_arr[i]->Mask,
            (ace_arr[i]->Mask & READ_PERMISSIONS) ? 'R' : '-',
            (ace_arr[i]->Mask & WRITE_PERMISSIONS) ? 'W' : '-',
            (ace_arr[i]->Mask & EXECUTE_PERMISSIONS) ? 'X' : '-');
        
        // print the SID
        if(!ConvertSidToStringSid((PSID)&ace_arr[i]->SidStart, &sid_str))
        {
            printf("\n");
            print_error(L"ConvertSidToString", GetLastError());
            continue;
        }
        
        wprintf(L"%s\n", sid_str);
        LocalFree(sid_str);
    }
}

/*
 * Moves the ACE at index src to index dst
 */
BOOL move_ace(ACCESS_ALLOWED_ACE **ace_arr, int len, int src, int dst)
{
    int i;
    ACCESS_ALLOWED_ACE *temp = NULL;
    
    if((src < 0 || src >= len) || (dst < 0 || dst >= len))
    {
        printf("ERROR: Index out of range.\n");
        return FALSE;
    }

    temp = ace_arr[src];

    if(dst < src)
    {
        // shift the entries down the list to make room for ace
        for(i = src; i > dst; i--)
            ace_arr[i] = ace_arr[i - 1];
    }
    else if(dst > src)
    {
        // shift the entries up the list to make room for ace
        for(i = src; i < dst; i++)
            ace_arr[i] = ace_arr[i + 1];
    }

    ace_arr[dst] = temp;
    return TRUE;
}

/*
 * Reorders the ACEs in the list specified by the order in the index array
 */
BOOL reorder_acl(ACCESS_ALLOWED_ACE **ace_arr, int len, int *idx_order, int idx_len)
{
    int i, j;
    ACCESS_ALLOWED_ACE *temp = NULL;
    ACCESS_ALLOWED_ACE **new_arr = NULL;

    // check all the indicies that they are in range
    for(i = 0; i < idx_len; i++)
    {
        if(idx_order[i] < 0 || idx_order[i] >= len)
        {
            printf("ERROR: Index out of range.\n");
            return FALSE;
        }
    }

    // check for duplicate entries (the slow way)
    for(i = 0; i < idx_len; i++)
    {
        for(j = 0; j < idx_len; j++)
        {
            if(i != j && idx_order[i] == idx_order[j])
            {
                printf("ERROR: Duplicate index.\n");
                return FALSE;
            }
        }
    }

    // allocate a new temporary array to use
    if((new_arr = (ACCESS_ALLOWED_ACE **)calloc(len, sizeof(ACCESS_ALLOWED_ACE *))) == NULL)
    {
        printf("ERROR: calloc() failed.\n");
        return FALSE;
    }
    
    // copy entries specified in the order list to the new array
    for(i = 0; i < idx_len; i++)
    {
        new_arr[i] = ace_arr[idx_order[i]];
        ace_arr[idx_order[i]] = NULL;
    }

    // copy the rest of the entries to the end of the new array, in order
    for(j = 0, i = idx_len; j < len; j++)
    {
        if(ace_arr[j] != NULL)
        {
            new_arr[i++] = ace_arr[j];
            ace_arr[j] = NULL;
        }
    }

    // copy entries back to original array;
    for(i = 0; i < len; i++)
        ace_arr[i] = new_arr[i];

    free(new_arr);
    return TRUE;
}

/*
 * Checks if a null-terminated string consists of only digits
 */
BOOL isnum(char *str)
{
    for(; *str != '\0'; str++)
    {
        if(!isdigit(*str))
            return FALSE;
    }

    return TRUE;
}

/*
 * Prints the error message with the function name
 */
void print_error(LPWSTR func_name, DWORD error_code)
{
    LPTSTR error_str = NULL;
    
    // gets the message string from the error code
    FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM, 0, GetLastError(), 0, (LPTSTR)&error_str, 0, NULL);
    wprintf(L"%s(%d): %s\n", func_name, error_code, error_str);

    if(error_str)
        LocalFree(error_str);
}

/*
 * Prints program usage.
 */
int usage(char *progname)
{
    printf("Reorders access control entries within a discretionary access control\n" \
           "list on NTFS files or directories.\n\n");
    printf("usage: %s <file or directory> [/M s d] [/R ...]\n", progname);
    printf("   /M s d     Move ACE at index s to index d\n"                             \
           "   /R i j...  Reorder ACEs at the specified indicies to the order given.\n" \
           "              If not all indicies are included, the missing ones will be\n" \
           "              appended to the end of the ACL in their respective order.\n"  \
           "              Example: For an ACL of length 4, \'/R 3 0 1\' will reorder\n" \
           "                       entries [0],[1],[2],[3] to [3],[0],[1],[2].\n");
    return -1;
}
